#include <opencv2/opencv.hpp>
#include <iostream>
#include <filesystem>
#include "line_detection.hpp"
#include "undistortion.hpp"

namespace fs = std::filesystem;

// 全局变量用于存储手动选择的直线
std::vector<cv::Vec4i> selected_lines;
// 全局变量用于存储当前显示的图像
cv::Mat current_display_image;

// 鼠标回调函数
void onMouse(int event, int x, int y, int flags, void* userdata) {
    if (event == cv::EVENT_LBUTTONDOWN) {
        std::vector<cv::Vec4i>* lines = static_cast<std::vector<cv::Vec4i>*>(userdata);
        cv::Point click_point(x, y);
        for (const auto& line : *lines) {
            cv::Point p1(line[0], line[1]);
            cv::Point p2(line[2], line[3]);
            // 计算点到直线的距离
            float dist = cv::pointPolygonTest(std::vector<cv::Point>{p1, p2}, click_point, true);
            if (std::abs(dist) < 5) { // 距离阈值可调整
                selected_lines.push_back(line);
                cv::line(current_display_image, p1, p2, cv::Scalar(0, 0, 255), 2);
                cv::imshow("Select lines", current_display_image);
                break;
            }
        }
    }
}

int main() {
    std::vector<std::string> fittingImages = {
        "../images/131325.jpg", "../images/131751.jpg",
        "../images/141035.jpg", "../images/132706.jpg", "../images/132930.jpg", "../images/142716.jpg"
    };

    std::vector<std::string> validationImages = {
        "../images/141142.jpg", "../images/125922.jpg", "../images/125921.jpg", "../images/131752.jpg","../images/141036.jpg", "../images/135852.jpg", "../images/141143.jpg"
    };

    std::vector<cv::Point2f> all_sampled_points;

    for (const auto& path : fittingImages) {
        std::cout << "处理拟合图像: " << path << std::endl;

        cv::Mat img = cv::imread(path);
        if (img.empty()) {
            std::cerr << "无法读取图像: " << path << std::endl;
            continue;
        }

        cv::Mat gray, blurred, edges;
        cv::cvtColor(img, gray, cv::COLOR_BGR2GRAY);
        cv::GaussianBlur(gray, blurred, cv::Size(5, 5), 1.5);
        cv::Canny(blurred, edges, 50, 150);

        std::vector<cv::Vec4i> lines;
        cv::HoughLinesP(edges, lines, 1, CV_PI / 180, 80, 30, 10);

        // 绘制检测到的直线
        current_display_image = img.clone();
        for (const auto& line : lines) {
            cv::line(current_display_image, cv::Point(line[0], line[1]), cv::Point(line[2], line[3]), cv::Scalar(0, 255, 0), 2);
        }

        // 创建可调整大小的窗口
        cv::namedWindow("Select lines", cv::WINDOW_NORMAL);
        // 调整窗口大小以适应图像
        cv::resizeWindow("Select lines", img.cols, img.rows);

        // 显示图像并设置鼠标回调函数
        cv::imshow("Select lines", current_display_image);
        cv::setMouseCallback("Select lines", onMouse, &lines);

        std::cout << "请点击需要保留的直线，按任意键结束选择。" << std::endl;
        cv::waitKey(0);
        cv::destroyWindow("Select lines");

        cv::Point2f center(img.cols / 2.0f, img.rows / 2.0f);

        for (auto& l : selected_lines) {
            auto raw_pts = sampleLinePoints(
                cv::Point2f(static_cast<float>(l[0]), static_cast<float>(l[1])),
                cv::Point2f(static_cast<float>(l[2]), static_cast<float>(l[3])),
                30
            );
            for (auto& pt : raw_pts) {
                float x = (pt.x - center.x) / center.x;
                float y = (pt.y - center.y) / center.y;
                all_sampled_points.emplace_back(x, y);
            }
        }

        // 清空选择的直线
        selected_lines.clear();
    }

    if (all_sampled_points.size() < 50) {
        std::cerr << "检测到的直线点过少，无法拟合。" << std::endl;
        return -1;
    }

    auto [best_k1, best_k2] = optimize_k1_k2(all_sampled_points);
    std::cout << "拟合完成，最佳参数: k1 = " << best_k1 << ", k2 = " << best_k2 << std::endl;

    fs::create_directories("results");

    // 对拟合图像进行畸变校正并保存
    for (const auto& path : fittingImages) {
        cv::Mat img = cv::imread(path);
        if (img.empty()) {
            std::cerr << "无法读取拟合图像: " << path << std::endl;
            continue;
        }

        cv::Mat undistorted;
        undistort_image(img, undistorted, best_k1, best_k2);

        std::string filename = fs::path(path).filename().stem().string();
        std::string ext = fs::path(path).extension().string();
        std::string out_path = "../results/" + filename + "_undistorted" + ext;

        cv::imwrite(out_path, undistorted);
        std::cout << "保存校正拟合图像: " << out_path << std::endl;
    }

    // 对验证图像进行畸变校正并保存
    for (const auto& path : validationImages) {
        cv::Mat img = cv::imread(path);
        if (img.empty()) {
            std::cerr << "无法读取验证图像: " << path << std::endl;
            continue;
        }

        cv::Mat undistorted;
        undistort_image(img, undistorted, best_k1, best_k2);

        std::string filename = fs::path(path).filename().stem().string();
        std::string ext = fs::path(path).extension().string();
        std::string out_path = "../results/" + filename + "_undistorted" + ext;

        cv::imwrite(out_path, undistorted);
        std::cout << "保存校正验证图像: " << out_path << std::endl;
    }

    std::cout << "✅ 全部处理完成！" << std::endl;
    return 0;
}

